import os
from PIL import Image
import numpy as np
from segment_anything import SamAutomaticMaskGenerator, sam_model_registry
import matplotlib.pyplot as plt
import torch
import random
from tqdm import tqdm
import argparse
import warnings
warnings.filterwarnings('ignore', category=UserWarning)

from utils.common import check_and_create_directory, list_all_files, chunk_list
from utils.auto_sam import auto_sam
from utils.vit.vit import vit_encode


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Process some paths.')
    parser.add_argument('--mode', type=str, required=True, help='Mode (train/val)')
    args = parser.parse_args()

    mode = args.mode
    input_path = '/home/zry/datasets/building/%s/rgbs' % mode
    output_path = '/home/zry/datasets/building/%s/seg' % mode
    tmp_path = '/home/zry/datasets/building/%s/tmp' % mode
    
    batch_size = 64
    device = "cuda:0"

    check_and_create_directory(output_path)
    check_and_create_directory(tmp_path)

    image_paths = list_all_files(input_path, '.jpg')

    debug_idx = 1257
    check_idx = len(list_all_files(tmp_path, '.pt'))

    # 先直接全部输进去吧，不行就分batch
    for idx, image_path in tqdm(enumerate(image_paths), total=len(image_paths), desc='image segementation'):
        # if idx < check_idx: continue
        if idx != debug_idx: continue
        try:
            image_label = os.path.split(image_path)[-1] 
            masks = auto_sam(image_path, tmp_path, device)
            # masks = [{"segmentation": x['segmentation']} for x in masks]  # 只保留其中的分割结果
            
            new_pt = torch.ones(masks[0]['segmentation'].shape) * -1
            for idx, mask in tqdm(enumerate(masks), total=len(masks), desc='masks => one image', leave=False):
                seg_mask = mask['segmentation']
                new_pt[seg_mask] = idx
                
            new_f = os.path.join(tmp_path, f"%s.pt" % image_label)
            torch.save(new_pt, new_f)
        except:
            print('error', idx)
            pass
    
    # torch.save(tmp_path, 'seg_images')

    print('You can contine run `python vit.main.py`.')
    print('Note!!! You need to check all .pt files')
   






    
